#!/bin/bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m' # No Color

# 错误处理函数
handle_error() {
    echo -e "${RED}错误: $1${NC}"
    exit 1
}

# 提示信息函数
info() {
    echo -e "${GREEN}信息: $1${NC}"
}

# 警告信息函数
warning() {
    echo -e "${YELLOW}警告: $1${NC}"
}

# 备份当前源
backup_sources() {
    local backup_dir="/etc/apt/sources.list.d/backup_$(date +%s)"
    info "正在备份当前源配置..."
    mkdir -p "$backup_dir" || handle_error "创建备份目录失败"
    
    if [ -f "/etc/apt/sources.list" ]; then
        cp "/etc/apt/sources.list" "$backup_dir/sources.list.bak" || handle_error "备份sources.list失败"
        info "已备份到: $backup_dir/sources.list.bak"
    else
        warning "未找到现有sources.list文件，将创建新文件"
    fi
}

# 设置阿里云镜像源
set_aliyun_mirror() {
    info "正在设置阿里云镜像源..."
    echo "deb http://mirrors.aliyun.com/kali kali-rolling main non-free contrib
deb-src http://mirrors.aliyun.com/kali kali-rolling main non-free contrib" | sudo tee /etc/apt/sources.list >/dev/null || handle_error "写入源配置失败"
    info "已设置为阿里云镜像源"
}

# 更新源
update_sources() {
    info "正在更新软件包索引..."
    sudo apt update || handle_error "更新软件包索引失败"
    info "软件包索引更新完成"
}

# 检查是否为root用户或使用sudo
if [ "$(id -u)" -ne 0 ]; then
    handle_error "请使用sudo权限运行此脚本"
fi

# 显示欢迎信息
echo "========================================="
echo "      Kali Linux 阿里云镜像源设置工具"
echo "========================================="

# 检查网络连接
ping -c 1 www.baidu.com >/dev/null 2>&1
if [ $? -ne 0 ]; then
    warning "网络连接可能存在问题，某些操作可能会失败"
fi

# 执行切换操作
backup_sources
set_aliyun_mirror
update_sources

# 完成提示
echo "========================================="
info "已成功将系统源切换为阿里云镜像源！"
echo "========================================="
exit 0

